##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggsci)

##########################################################################################

option_list <- list(
    make_option(c("--gene"), type = "character") ,
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--ccf_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    gene <- "MUC6"
    info_file <- "~/20220915_gastric_multiple/dna_combinePublic/baseTable/STAD_Info.addBurden.MSI_MSS.tsv"
    ccf_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.tsv"
	images_path <- paste(work_dir,"/images/mutBaselinePlot/CFTR",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

gene <- opt$gene
info_file <- opt$info_file
ccf_file <- opt$ccf_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")
col <- col[1:3]

col_im <- brewer.pal(9,"YlGnBu")[6:8]
names(col_im) <- c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)")

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
## 纳入考虑的宏观变量
base_col <- c("Gender" , "Age_divide" , "Tobacco" , "Alcohol" , "PickleFood" , "HP")

###########################################################################################

dat_ccf <- data.frame(fread( ccf_file ))
dat_info <- data.frame(fread( info_file ))

###########################################################################################
## 去除MSI
dat_info <- subset( dat_info , MS_Type != "MSI" & Type != "IM + IGC + DGC" )

###########################################################################################

dat_ccf <- subset(dat_ccf , Hugo_Symbol==gene & Variant_Classification %in% Variant_Types)

mutPatient_IM <- unique(dat_info[dat_info$Tumor %in% dat_ccf$Sample & dat_info$Class=="IM","Patient"])
mutPatient_IGC <- unique(dat_info[dat_info$Tumor %in% dat_ccf$Sample & dat_info$Class=="IGC","Patient"])
mutPatient_DGC <- unique(dat_info[dat_info$Tumor %in% dat_ccf$Sample & dat_info$Class=="DGC","Patient"])

dat_info$DriverMut_IM <- ifelse( dat_info$Patient %in% mutPatient_IM , "Mut" , "NoMut" )
dat_info$DriverMut_IGC <- ifelse( dat_info$Patient %in% mutPatient_IGC , "Mut" , "NoMut" )
dat_info$DriverMut_DGC <- ifelse( dat_info$Patient %in% mutPatient_DGC , "Mut" , "NoMut" )

###########################################################################################
## 突变合并
dat_plot_IM <- subset( dat_info , Class == "IM" )
dat_plot_IM <- dat_plot_IM %>%
group_by( Patient , Age , Gender , Tobacco , Alcohol , PickleFood , HP , Type ) %>%
summarize( DriverMut = unique(DriverMut_IM))
dat_plot_IM$Class <- "IM"

dat_plot_IGC <- subset( dat_info , Class == "IGC" )
dat_plot_IGC <- dat_plot_IGC %>%
group_by( Patient , Age , Gender , Tobacco , Alcohol , PickleFood , HP , Type ) %>%
summarize( DriverMut = unique(DriverMut_IGC))
dat_plot_IGC$Class <- "IGC"

dat_plot_DGC <- subset( dat_info , Class == "DGC" )
dat_plot_DGC <- dat_plot_DGC %>%
group_by( Patient , Age , Gender , Tobacco , Alcohol , PickleFood , HP , Type ) %>%
summarize( DriverMut = unique(DriverMut_DGC))
dat_plot_DGC$Class <- "DGC"

dat_plot <- rbind( dat_plot_IM , dat_plot_IGC , dat_plot_DGC )
dat_plot$Class <- factor( dat_plot$Class , levels = c("IGC" , "DGC" , "IM") , order = T )
dat_plot <- data.frame(dat_plot)

###########################################################################################

trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("p == ",down," %*% 10","^",up)
    return(text)
}

###########################################################################################
## 年龄按中位数分
median_age <- median(unique(data.frame(dat_plot$Patient , dat_plot$Age))$dat_plot.Age , na.rm = T)
dat_plot$Age_divide <- ifelse(dat_plot$Age > median_age , "Older" , "Younger")

###########################################################################################
## 计算P值
plotMutRate <- function(dat_plot_tmp = dat_plot_tmp , baseuse = baseuse , class_type = class_type , col_use = col_use){

    dat_plot_tmp$p.value = ""
    dat_plot_tmp$p_text = ""

    result <- c()
    for(classN in unique(dat_plot_tmp$Class)){

        print(classN)

        tmp <- subset( dat_plot_tmp , Class == classN )
        tmp_data <- data.frame(table(tmp$DriverMut , tmp$useCol))
        if( length(unique(tmp$useCol)) > 1 & length(unique(tmp$DriverMut)) > 1 ){
            p <- fisher.test(table(tmp$DriverMut , tmp$useCol))$p.value
        }else{
            p <- 1
        }

        colnames(tmp_data) <- c("DriverMut" , "Baseline" , "SampleNum")
        tmp_data <- tmp_data %>%
        group_by( DriverMut ) %>%
        summarize( Baseline = Baseline , SampleNum = SampleNum , Ratio=SampleNum/(sum(SampleNum)) )

        tmp_data$p.value <- p
        tmp_data$Class <- classN
        result <- rbind(result , tmp_data)
    }

    if( p < 0.001 ){
        p_text <- trans(result$p.value)
    }else{
        p_text <- paste0( "p == " , round(as.numeric(result$p.value) , 3) ) 
    }

    result$p_text <- p_text
    result$value_text <- paste0( round(result$Ratio , 2) * 100 , "%" , "\n" , result$SampleNum)
    result$Class <- factor( result$Class , levels = names(col_use) , order = T )
    result[result$SampleNum==0,"value_text"] <- ""

    plot <- ggplot( data = result , aes( x = DriverMut , y = Ratio , fill = Baseline ))+
    geom_bar(position = "stack", stat = "identity") + 
    theme_bw()+
    labs(x="",y="Mutation Rate")+
    facet_grid(.~Class) +
    theme(panel.grid = element_blank())+
    scale_fill_npg() +
    ylim(0,1.05)+
    geom_text(aes(label=p_text , y = 1.05 ,x = 1.5),parse = TRUE,size=4)+
    geom_text(aes(label=value_text) , position=position_stack(vjust = 0.5) , size=3 , color="black")+
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
                legend.position ='right',
                legend.title = element_blank() ,
                panel.grid.major=element_line(colour=NA),
                plot.title = element_text(size = 12,color="black",face='bold'),
                legend.text = element_text(size = 12,color="black",face='bold'),
                axis.text.y = element_text(size = 7,color="black",face='bold'),
                axis.title.x = element_text(size = 12,color="black",face='bold'),
                axis.title.y = element_text(size = 12,color="black",face='bold'),
                strip.text.x = element_text(size = 7 , face = 'bold'),
                axis.ticks.x = element_blank(),
                axis.text.x = element_text(size = 8,color="black",face='bold') ,
                axis.line = element_line(size = 0.5))

    out_name <- paste0( images_path , "/",baseuse,".CompareRate.",class_type,".pdf" ) 
    ggsave( out_name , plot , width = 5 , height = 5 )

    out_name <- paste0( images_path , "/",baseuse,".CompareRate.",class_type,".tsv" )  
    write.table( result , out_name , row.names = F , quote = F , sep = "\t" )
}

###########################################################################################
class_type <- "IM_IGC_DGC"
for( baseuse in base_col ){

    dat_plot_tmp <- data.frame(dat_plot)
    dat_plot_tmp$useCol <- dat_plot[[baseuse]]
    dat_plot_tmp <- subset( dat_plot_tmp , !is.na(useCol) )

    if(baseuse=="HP"){
        dat_plot_tmp[dat_plot_tmp$useCol=="Negative","useCol"] <- "HP_N"
        dat_plot_tmp[dat_plot_tmp$useCol=="Positive","useCol"] <- "HP_Y"
        dat_plot_tmp$useCol <- factor( dat_plot_tmp$useCol , levels = c("HP_Y" , "HP_N") , order = T )
    }

    col <- col[c("IGC" , "DGC" , "IM")]
    
    plotMutRate(dat_plot_tmp = dat_plot_tmp , baseuse = baseuse , class_type = class_type , col_use = col)
}

###########################################################################################
## 不同胃癌亚型的肠化
dat_plot_tmp <- subset( dat_plot , Class == "IM" )
dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + IGC" , "IM(IGC)" , dat_plot_tmp$Class )
dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + DGC" , "IM(DGC)" , dat_plot_tmp$Class )
dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + IGC + DGC" , "IM(IGC_DGC)" , dat_plot_tmp$Class )
dat_plot_tmp$Class <- factor( dat_plot_tmp$Class , levels = c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)") , order = T )
dat_plot <- dat_plot_tmp
class_type <- "IM"

for( baseuse in base_col ){
    print(baseuse)

    dat_plot_tmp <- data.frame(dat_plot)
    dat_plot_tmp$useCol <- dat_plot[[baseuse]]
    dat_plot_tmp <- subset( dat_plot_tmp , !is.na(useCol) )

    if(baseuse=="HP"){
        dat_plot_tmp[dat_plot_tmp$useCol=="Negative","useCol"] <- "HP_N"
        dat_plot_tmp[dat_plot_tmp$useCol=="Positive","useCol"] <- "HP_Y"
        dat_plot_tmp$useCol <- factor( dat_plot_tmp$useCol , levels = c("HP_Y" , "HP_N") , order = T )
    }
    
    plotMutRate(dat_plot_tmp = dat_plot_tmp , baseuse = baseuse , class_type = class_type , col_use = col_im)
}
